

class FoodDataset(Dataset):
    def __init__(self, image_path, image_size=(128, 128), mode='train'):
        self.image_path = image_path
        self.image_file_list = sorted(os.listdir(image_path))
        self.mode = mode
        # training 时做 data augmentation
        self.train_transforms = Compose([
            Resize(size=image_size),
            RandomHorizontalFlip(),
            RandomRotation(15),
            Transpose(),
            Normalize(mean=127.5, std=127.5)
        ])
        # testing 时不需做 data augmentation
        self.test_transforms = Compose([
            Resize(size=image_size),
            Transpose(),
            Normalize(mean=127.5, std=127.5)
        ])
        
    def __len__(self):
        return len(self.image_file_list)
    
    def __getitem__(self, idx):
        img = cv2.imread(os.path.join(self.image_path, self.image_file_list[idx]))
        if self.mode == 'train':
            img = self.train_transforms(img)
            label = int(self.image_file_list[idx].split("_")[0])
            return img, label
        else:
            img = self.test_transforms(img)
            return img